import imp
import math

import numpy as np
import tensorflow as tf
import t3f


flags = tf.app.flags
FLAGS = flags.FLAGS

# 模型选择
flags.DEFINE_string('flag_net_module', './TrainingMnist.py', 'Module selection with specific dataset.')


def run_vars_and_ops(sess, l_tts, l_ops, dict_feeder):
	pre_loss = 0.0
	while True:
		loss, _ = sess.run(l_ops, dict_feeder)
		diff = math.fabs(loss - pre_loss)
		pre_loss = loss
		if diff < 1e-4:
			break

	l_tt_data = sess.run(l_tts)
	return l_tt_data


def main(_):
	ds = imp.load_source('module', FLAGS.flag_net_module)
	dict_dataset, dict_mean_std = ds.interface_get_dataset(False)

	# 待TT化的输入数据placeholder
	tfph_input_data = tf.placeholder(dtype = tf.float32, shape = dict_mean_std['mean']['mean'][0].shape, name = 'ph_input_data')

	# 输入数据TT化的计算
	l_tts, l_ops = ds.interface_riemannian_data(tfph_input_data)

	sess = tf.Session(config = tf.ConfigProto(allow_soft_placement = True))
	sess.run(tf.global_variables_initializer())

	# 训练数据TT化
	n_train_size = dict_dataset['train']['train_labels'].shape[0]
	for i in range(n_train_size):
		dict_feeder = {tfph_input_data : dict_dataset['train']['train_data'][0][i]}
		l_tt_data = run_vars_and_ops(sess, l_tts, l_ops, dict_feeder)
		ds.interface_save_tt_data(l_tt_data, True)
		print('The %d/%d sample (train) has done.' % ((i + 1), n_train_size))

	# 验证数据TT化
	n_validation_size = dict_dataset['validation']['validation_labels'].shape[0]
	for i in range(n_validation_size):
		dict_feeder = {tfph_input_data : dict_dataset['validation']['validation_data'][0][i]}
		l_tt_data = run_vars_and_ops(sess, l_tts, l_ops, dict_feeder)
		ds.interface_save_tt_data(l_tt_data, False)
		print('The %d/%d sample (validation) has done.' % ((i + 1), n_validation_size))

	# mean与std数据TT化并保存
	dict_feeder = {tfph_input_data : dict_mean_std['mean']['mean'][0]}
	l_tt_mean = run_vars_and_ops(sess, l_tts, l_ops, dict_feeder)
	dict_feeder = {tfph_input_data : dict_mean_std['std']['std'][0]}
	l_tt_std = run_vars_and_ops(sess, l_tts, l_ops, dict_feeder)
	ds.interface_save_tt_dataset(l_tt_mean, l_tt_std, dict_dataset['train']['train_labels'], dict_dataset['validation']['validation_labels'])
	print('Dataset in TT format is over.')


if __name__ == '__main__':
    tf.app.run()
